package com.wave.wavesystem.ai.test;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class MyMultiQueryExpender {


    private final ChatClient.Builder clientBuilder;

    public MyMultiQueryExpender(ChatModel dashscopeChatClient) {
        this.clientBuilder = ChatClient.builder(dashscopeChatClient);
    }


    public List<Query> getQueries(String q) {
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(clientBuilder)
                .numberOfQueries(3)
                .build();
        List<Query> queries = queryExpander.expand(new Query(q));
        return queries;
    }


}
